package ex3.visitors;

import common.exceptions.SemanticError;
import ic.ast.Visitor;
import ic.ast.decl.*;
import ic.ast.expr.*;
import ic.ast.expr.BinaryOp.BinaryOps;
import ic.ast.expr.UnaryOp.UnaryOps;
import ic.ast.stmt.*;
import ic.symbols.*;
import ic.types.MethodType;
import ic.types.TypeTable;

import java.util.Iterator;

/**
 * Created with IntelliJ IDEA.
 * User: Roi
 * Date: 12/6/13
 *
 * A visitor responsible for type and scope checks
 */
public class TypeCheckVisitor implements Visitor {

	// for checking that 'this' is used in the correct context.
	// the flag is turned on when entering a virtual method
	// and off when exiting it
	private boolean inVirtualMethodContext = false;

    @Override
	public Object visit(Program program) throws Exception {
		for (DeclClass icClass : program.getClasses()) {
			icClass.accept(this);
		}
        return null;
	}

	@Override
	public Object visit(DeclClass icClass) throws Exception {
		for (DeclField field : icClass.getFields()) {
			field.accept(this);
		}
		for (DeclMethod method : icClass.getMethods()) {
			method.accept(this);
		}
		return null;
	}

	@Override
	public Object visit(DeclField field) throws Exception {
		return TypeTable.getType(field.getType(), false);
	}

	@Override
	public Object visit(DeclVirtualMethod method) throws Exception {
		inVirtualMethodContext = true;
		for (Statement stmt : method.getStatements()) {
			stmt.accept(this);
		}
		inVirtualMethodContext = false;
		return null;
	}

	@Override
	public Object visit(DeclStaticMethod method) throws Exception {
		for (Statement stmt : method.getStatements()) {
			stmt.accept(this);
		}
		return null;
	}

	@Override
	public Object visit(DeclLibraryMethod method) throws Exception {

		//		for (Parameter param : method.getFormals()) {
		//			param.accept(this);
		//		}
		return TypeTable.getType(method);
	}

	@Override
	public Object visit(Parameter formal) throws Exception {
		return TypeTable.getType(formal.getType(), false);
	}

	@Override
	public Object visit(PrimitiveType type) throws Exception {
		return TypeTable.getType(type, false);
	}

	@Override
	public Object visit(ClassType type) throws Exception {
		try {
			TypeTable.getClassTypeByName(type.getDisplayName());
		} catch (SemanticError semantic) {
			typeError(type.getLine(), semantic.getMessage());
		}
		return TypeTable.getType(type);
	}

    /**
     * Check:
     * Assignment type is a subtype of the variable assigned to.
     * Returns null, or an error if found.
     */
    @Override
	public Object visit(StmtAssignment assignment) throws Exception {

        Type variableType = (Type) assignment.getVariable().accept(this);
        if (variableType == null)
            return null;

		Type assignmentType = (Type) assignment.getAssignment().accept(this);
		if (assignmentType == null)
			return null;

		if (!(assignmentType.subTypeOf(variableType))) {
			typeError(assignment.getLine(), "Invalid assignment of type " + assignmentType.getOutputName()
                    + " to variable of type " + variableType.getOutputName());

		}
		return null;
	}

	@Override
	public Object visit(StmtCall callStatement) throws Exception {
		callStatement.getCall().accept(this);
		return null;
	}

    /**
     * Check: return value code is a subtype of the method type
     */
    @Override
	public Object visit(StmtReturn returnStatement) throws Exception {
		Type returnValueType;
		if (returnStatement.hasValue()) {
			returnValueType = (Type) returnStatement.getValue().accept(this);
		} else {
			returnValueType = TypeTable.getType("void");
		}
		String currentMethodName = returnStatement.getEnclosingScope().getCurrentMethodName();
		MethodSymbol currentMethodSymbol = (MethodSymbol) returnStatement.getEnclosingScope().lookup(currentMethodName);
		MethodType currentMethodType = currentMethodSymbol.getMetType();
		Type methodReturnType = TypeTable.getType(currentMethodType.getReturnType());
		if (!returnValueType.subTypeOf(methodReturnType)) {
			typeError(returnStatement.getLine(), "Return statement is not of type " + methodReturnType.getOutputName());
		}


		return null;
	}

    /**
     * Check: types in an If/Else statement.
     * Makes sure that the if condition is of type boolean.
     * Returns null, or an error if found.
     */
    @Override
	public Object visit(StmtIf ifStatement) throws Exception {
		Type conditionType = (Type) ifStatement.getCondition().accept(this);
		if (conditionType != TypeTable.getType("boolean")) {
			typeError(ifStatement.getLine(), "Non boolean condition for if statement");
		}

		ifStatement.getOperation().accept(this);

		if (ifStatement.hasElse()) {
			ifStatement.getElseOperation().accept(this);
		}

		return null;
	}

    /**
     * Check: while statement condition is of type boolean.
     * Returns null, or an error if found.
     */
    @Override
	public Object visit(StmtWhile whileStatement) throws Exception {
		Type conditionType = (Type) whileStatement.getCondition().accept(this);
		if (conditionType != TypeTable.getType("boolean")) {
			typeError(whileStatement.getLine(), "Non boolean condition for while statement");
		}

		whileStatement.getOperation().accept(this);
		return null;
	}

	@Override
	public Object visit(StmtBreak breakStatement) {
		return null;
	}

	@Override
	public Object visit(StmtContinue continueStatement) {
		return null;
	}

	@Override
	public Object visit(StmtBlock statementsBlock) throws Exception {
		for (Statement stmt : statementsBlock.getStatements()) {
			stmt.accept(this);
		}
		return null;
	}

    /**
     * Check: variable type is a subtype of the local variables type
     */
    @Override
	public Object visit(LocalVariable localVariable) throws Exception {
		localVariable.getType().accept(this);
		Type varType = TypeTable.getType(localVariable.getType(), false);

		if (localVariable.getInitialValue() != null) {
			Type initValueType = (Type) localVariable.getInitialValue().accept(this);
			if (!initValueType.subTypeOf(varType)) {
				typeError(localVariable.getLine(), "type mismatch: expected " + varType + " and got " + initValueType);
			}
		}

		return null;
	}

    /**
     * Return location type, or an error if found.
     */
    @Override
	public Object visit(RefVariable location) throws Exception {
		Type expressionType = null;
		if (TypeTable.getType(location.getName()) == null) {
			Symbol varSymbol = location.getEnclosingScope().lookup(location.getName());
			if (varSymbol == null) {
				typeError(location.getLine(), location.getName() + " not found in symbol table");
			}
			if (varSymbol.getKind() == Kind.VARIABLE) {
				if (varSymbol.getLine() > location.getLine() ||
						(varSymbol.getLine() == location.getLine())) {
					// the variable we're looking for was defined in an enclosing scope or it was
					// used before declaration
					Symbol altSymbol = location.getEnclosingScope().getParent().lookup(location.getName());
					if (altSymbol != null && (altSymbol.getKind() == Kind.VARIABLE || altSymbol.getKind() == Kind.FIELD)) {
						varSymbol = altSymbol;
					} else {
						typeError(location.getLine(), "variable " + location.getName() + " is used before declaration");
					}
				}
			}
			if (varSymbol.getKind() == Kind.VARIABLE) {
				expressionType = TypeTable.getType(((VarSymbol) varSymbol).getType(), false);
			} else if (varSymbol.getKind() == Kind.FIELD && inVirtualMethodContext) {
				expressionType = TypeTable.getType(((FieldSymbol) varSymbol).getType(), false);
			} else {
				typeError(location.getLine(), "Use of field inside static method is not allowed");
			}
		} else {
			Type classType = TypeTable.getType(location.getName());
			Symbol fieldSymbol = null;
			try {
				SymbolTable classSymbolTable = TypeTable.getClassTypeByName(classType.getDisplayName()).getEnclosingScope().getSymbolTable(classType.getDisplayName());
				fieldSymbol = classSymbolTable.lookup(location.getName());
			} catch (SemanticError semantic) {
				typeError(location.getLine(), "no such class: " + classType.getDisplayName());
			}
			if (fieldSymbol == null) {
				typeError(location.getLine(), "undeclared identifier: " + location.getName());
			}
			if (fieldSymbol.getKind() != Kind.FIELD) {
				typeError(location.getLine(), "illegal reference: " + location.getName());
			}
			expressionType = TypeTable.getType(((FieldSymbol) fieldSymbol).getType(), false);
		}
		return expressionType;
	}


	@Override
	public Object visit(RefField location) throws Exception {
		Type expressionType = null;
		if (location.getObject() == null) {
			Symbol varSymbol = location.getEnclosingScope().lookup(location.getField());
			if (varSymbol == null) {
				typeError(location.getLine(), "undeclared identifier: " + location.getField());
			}
			if (varSymbol.getKind() == Kind.VARIABLE) {
				if (varSymbol.getLine() > location.getLine() ||
						(varSymbol.getLine() == location.getLine())) {
					// the variable we're looking for was defined in an enclosing scope or it was
					// used before declaration
					Symbol altSymbol = location.getEnclosingScope().getParent().lookup(location.getField());
					if (altSymbol != null && (altSymbol.getKind() == Kind.VARIABLE || altSymbol.getKind() == Kind.FIELD)) {
						varSymbol = altSymbol;
					} else {
						typeError(location.getLine(), "variable " + location.getField() + " is used before declaration");
					}
				}
			}
			if (varSymbol.getKind() == Kind.VARIABLE) {
				expressionType = TypeTable.getType(((VarSymbol) varSymbol).getType(), false);
			} else if (varSymbol.getKind() == Kind.FIELD && inVirtualMethodContext) {
				expressionType = TypeTable.getType(((FieldSymbol) varSymbol).getType(), false);
			} else {
				typeError(location.getLine(), "illegal reference: " + location.getField());
			}
		} else {
			Type classType = (Type) location.getObject().accept(this);
			Symbol fieldSymbol = null;
			try {
				SymbolTable classSymbolTable = TypeTable.getClassTypeByName(classType.getDisplayName()).getEnclosingScope().getSymbolTable(classType.getDisplayName());
				fieldSymbol = classSymbolTable.lookup(location.getField());
			} catch (SemanticError semantic) {
				typeError(location.getLine(), "no such class: " + classType.getDisplayName());
			}
			if (fieldSymbol == null) {
				typeError(location.getLine(), "undeclared identifier: " + location.getField());
			}
			if (fieldSymbol.getKind() != Kind.FIELD) {
				typeError(location.getLine(), "illegal reference: " + location.getField());
			}
			expressionType = TypeTable.getType(((FieldSymbol) fieldSymbol).getType(), false);
		}
		return expressionType;
	}

    /**
     * Check: array index is of type int.
     * Return element type, or an error if found.
     */
    @Override
	public Object visit(RefArrayElement location) throws Exception {
		Type arrayType = (Type) location.getArray().accept(this);
		Type indexType = (Type) location.getIndex().accept(this);

		if (indexType != TypeTable.getType("int")) {
			typeError(location.getLine(), "array index must be of type int");
		}

		if (arrayType.getArrayDimension() < 1) {
			typeError(location.getLine(), arrayType + " is not an array type");
		}

		String typeName = arrayType.toString();
		String elementType = typeName.substring(0, typeName.length() - 2);

		return TypeTable.getType(elementType);

	}

    /**
     * Check:
     * types of arguments
     * method is defined in the correct class
     * all arguments correspond to the method's arguments types
     * Return method return value, or an error if found.
     */
    @Override
	public Object visit(StaticCall call) throws Exception {
		String className = call.getClassName();
		DeclClass classInstance = null;
		try {
			classInstance = TypeTable.getClassTypeByName(className);
		} catch (SemanticError e) {
			typeError(call.getLine(), className + ": no such class");
		}

		SymbolTable classScope = classInstance.getEnclosingScope().getSymbolTable(className);
		String staticMethodName = call.getMethod();

		Symbol methodSymbol = classScope.lookup(staticMethodName);

		if (methodSymbol == null || methodSymbol.getKind() != Kind.METHOD ||
				!((MethodSymbol) methodSymbol).isStatic()) {
			typeError(call.getLine(), "Method " + staticMethodName + " doesn't exist");
		}

		MethodType methodType = ((MethodSymbol) methodSymbol).getMetType();
		if (call.getArguments().size() != methodType.getParamTypes().size()) {
			typeError(call.getLine(), "Invalid number of arguments for " + className + "." + staticMethodName);
		}

		Iterator<Type> methodParams = methodType.getParamTypes().iterator();
		for (Expression arg : call.getArguments()) {
			Type paramType = TypeTable.getType(methodParams.next(), false);
			Type argType = (Type) arg.accept(this);
			if (!argType.subTypeOf(paramType)) {
				typeError(call.getLine(), "argument and parameter type mismatch in call to " + staticMethodName +
						". Expected " + paramType + " and got " + argType);
			}
		}

		return TypeTable.getType(methodType.getReturnType(), false);
	}

    /**
     * Check:
     * method is defined in the relevant class
     * all arguments correspond to the method's arguments types
     * Return method return value, or an error if found.
     */
    @Override
	public Object visit(VirtualCall call) throws Exception {
		DeclClass instanceClass = null;
		Type instanceType;
		SymbolTable classScope;
		String instanceClassName = null;

        String methodName = call.getMethod();
        Symbol methodSymbol = null;
        MethodType methodType = null;

		// first, determine the class context
		if (call.hasExplicitObject()) {
			Expression location = call.getObject();
			instanceType = (Type) location.accept(this);
			instanceClassName = instanceType.getDisplayName();
			try {
				instanceClass = TypeTable.getClassTypeByName(instanceClassName);
			} catch (SemanticError e) {
				typeError(call.getLine(), e.getMessage());
			}
			classScope = instanceClass.getEnclosingScope().getSymbolTable(instanceClassName);
            methodSymbol = classScope.lookup(methodName);
            methodType = ((MethodSymbol) methodSymbol).getMetType();
		} else {
			try {
				classScope = call.getEnclosingScope();
                methodSymbol = classScope.lookup(methodName);

                if (!call.hasExplicitObject() && !inVirtualMethodContext && !((MethodSymbol) methodSymbol).isStatic()) {
                    typeError(call.getLine(), "Calling a local virtual method from inside a static method is not allowed");
                }

                methodType = ((MethodSymbol) methodSymbol).getMetType();
                if (call.getArguments().size() != methodType.getParamTypes().size()) {
                    typeError(call.getLine(), "Invalid number of arguments for " + methodName);
                }

				instanceClass = TypeTable.getClassTypeByName(classScope.getName());
				instanceClassName = instanceClass.getName();
			} catch (SemanticError e) {
				typeError(call.getLine(), e.getMessage());
			}
		}

		if (methodSymbol.getKind() != Kind.METHOD) {
			typeError(call.getLine(), methodName + ": no such method in " + instanceClassName);
		}
		if (!call.hasExplicitObject() && !inVirtualMethodContext && !((MethodSymbol) methodSymbol).isStatic()) {
			typeError(call.getLine(), "Calling a local virtual method from inside a static method is not allowed");
		}

		Iterator<Type> methodParams = methodType.getParamTypes().iterator();
		for (Expression arg : call.getArguments()) {
			Type paramType = TypeTable.getType(methodParams.next(), false);
			Type argType = (Type) arg.accept(this);
			if (!argType.subTypeOf(paramType)) {
				typeError(call.getLine(), "argument and parameter type mismatch in call to " + methodName +
						". Expected " + paramType + " and got " + argType);
			}
		}
		return TypeTable.getType(methodType.getReturnType(), false);
	}

	@Override
	public Object visit(This thisExpression) throws Exception {
		return TypeTable.getType(new ClassType(thisExpression.getLine(), thisExpression.getEnclosingScope().getName()));
	}

    /**
     * Check: new class expression - class type is legal and exists.
     */
    @Override
	public Object visit(NewInstance newClass) throws Exception {
		String className = newClass.getName();
		Type classType = null;
		classType = TypeTable.getType(new ClassType(newClass.getLine(), className));
		return classType;
	}

    /**
     * Check: new array expression - element type is a legal type and that the size is of type int.
     * Return the array type.
	 */
	@Override
	public Object visit(NewArray newArray) throws Exception {
		newArray.getType().accept(this);
		Type tmpType = TypeTable.getType(newArray.getType(), false);
		Type arrayType = TypeTable.getType(tmpType.clone(tmpType.getArrayDimension() + 1));

		Type sizeType = (Type) newArray.getSize().accept(this);
		if (sizeType != TypeTable.getType("int")) {
			typeError(newArray.getLine(), "array size must be of type int");
		}

		return arrayType;
	}

	@Override
	public Object visit(Length length) throws Exception {
		Type arrType = (Type) length.getArray().accept(this);

		if (arrType.getArrayDimension() < 1) {
			typeError(length.getLine(), "length is applicable only for array types");
		}

		return TypeTable.getType("int");
	}

    /**
     * Check: types in math binary operations.
     * for + operation: gets either 2 ints or 2 strings. Return int or string.
     * for {/,*,-} operations, gets 2 ints and return int.
     * Send an error otherwise
     */
    @Override
	public Object visit(BinaryOp binaryOp) throws Exception {
		Type op1Type = (Type) binaryOp.getFirstOperand().accept(this);
		Type op2Type = (Type) binaryOp.getSecondOperand().accept(this);
		BinaryOps operator = binaryOp.getOperator();

		// If the operator is a logical operator
		if (operator == BinaryOps.LAND || operator == BinaryOps.LOR 
				|| operator == BinaryOps.GT || operator == BinaryOps.GTE 
				|| operator == BinaryOps.LT || operator == BinaryOps.LTE
				|| operator == BinaryOps.EQUAL || operator == BinaryOps.NEQUAL) {

			// Check for the || and && operators that both operands are boolean
			if ((operator == BinaryOps.LAND)|| (operator == BinaryOps.LOR)) {
				if (op1Type != TypeTable.getType("boolean") || op2Type != TypeTable.getType("boolean")){
					typeError(binaryOp.getLine(), "Logical "+ operator.getDescription() + 
					" is applicable only for boolean operands");
				}
			}
			//Check if one operand type is subtype of the other
			if (operator == BinaryOps.EQUAL || operator == BinaryOps.NEQUAL){
				if (!(op1Type.subTypeOf(op2Type)) && !(op2Type.subTypeOf(op1Type))){
					typeError(binaryOp.getLine(), 
							"comparison of foreign types with the operator: " + operator.getDescription());
				}
			}
			//
			if (operator == BinaryOps.LTE || operator == BinaryOps.LT 
					|| operator == BinaryOps.GTE || operator == BinaryOps.GT) {
				if (op1Type != TypeTable.getType("int") || op2Type != TypeTable.getType("int")) {
					typeError(binaryOp.getLine(), "Invalid logical binary op (" + operator + ") on non-integer expression");
				}
			}
			return TypeTable.getType("boolean");
		}

		if ((op1Type == null) || (op2Type == null)) return null;

		if ((op1Type != TypeTable.getType("null")) || (op2Type != TypeTable.getType("null"))) {
			if (op2Type == TypeTable.getType("null"))
				op2Type = TypeTable.getType("string");

			if (op1Type == TypeTable.getType("null"))
				op1Type = TypeTable.getType("string");
		}

		if (op1Type != op2Type)
			typeError(binaryOp.getLine(), "Illegal " + binaryOp.getOperator().toString() +
					" operation. Both operands must be of the same type " +
					(binaryOp.getOperator() == BinaryOp.BinaryOps.PLUS ? "(int or string)" : "(int)"));

		if (binaryOp.getOperator() != BinaryOp.BinaryOps.PLUS) {
			if (op1Type != TypeTable.getType("int")) {
				typeError(binaryOp.getLine(), binaryOp.getOperator().toString() + " operations are only for two operands of type int");
			}
		} else {
			if (op1Type != TypeTable.getType("int") && op1Type != TypeTable.getType("string"))
				typeError(binaryOp.getLine(), binaryOp.getOperator().toString() + " operation is only for int or string operands");
		}
		return TypeTable.getType(op1Type, false);
	}

	/**
	 * Return the type of the literal
	 */
	@Override
	public Object visit(Literal literal) {
		switch (literal.getType()) {
            case STRING:
                return TypeTable.getType("string");
            case INT:
                return TypeTable.getType("int");
            case BOOLEAN:
                return TypeTable.getType("boolean");
            case VOID:
                return TypeTable.getType("null");
            }
		return null;
	}

	public Object visit(UnaryOp unaryOp) throws Exception {
		Type opType = (Type) unaryOp.getOperand().accept(this);
		UnaryOps operator = unaryOp.getOperator();

		if (operator == UnaryOps.LNEG) {
			if (opType != TypeTable.getType("boolean")) {
				typeError(unaryOp.getLine(), unaryOp.getOperator().toString() + " with a non-boolean operand");
			}			
			return TypeTable.getType("boolean");
		}
		else {
			if (opType != TypeTable.getType("int"))
				typeError(unaryOp.getLine(), "Unary "+ unaryOp.getOperator().toString() + " with a non-int operand");
			return TypeTable.getType("int");
		}

	}

	@Override
	public Object visit(GlobalSymbolTable table) {
		return null;
	}

	@Override
	public Object visit(ClassSymbolTable table) {
		return null;
	}

	@Override
	public Object visit(BlockSymbolTable table) {
		return null;
	}

	@Override
	public Object visit(MethodSymbolTable table) {
		return null;
	}

	@Override
	public Object visit(SymbolTable symbolTable) {
		return null;
	}

	private void typeError(int line, String message) throws Exception {
		throw new SemanticError(line, message);
	}
}
